/**
 * File: VelvetExec.java
 * Created by: mhaimel
 * Created on: 23 Jul 2009
 * CVS:  $Id: VelvetExec.java,v 1.2 2009/10/09 15:33:58 mhaimel Exp $
 */
package uk.ac.ebi.curtain.util;

import java.io.File;
import java.io.IOException;
import java.io.OutputStream;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collection;
import java.util.Collections;
import java.util.HashMap;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Map.Entry;
import java.util.Set;
import java.util.SortedMap;
import java.util.TreeMap;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.atomic.AtomicLong;

import org.apache.commons.io.FileUtils;
import org.apache.commons.io.IOUtils;
import org.apache.commons.lang.StringUtils;
import org.apache.commons.logging.Log;
import org.apache.commons.logging.LogFactory;

import uk.ac.ebi.curtain.utils.CurtainUncheckedException;
import uk.ac.ebi.curtain.utils.data.FileType;
import uk.ac.ebi.curtain.utils.data.ReadType;
import uk.ac.ebi.curtain.utils.file.FileInfo;
import uk.ac.ebi.curtain.utils.io.impl.FileIO;
import uk.ac.ebi.velvet.VelvetUncheckedException;
import uk.ac.ebi.velvet.exec.RuntimeExec;
import uk.ac.ebi.velvet.processor.StatisticsInfo;
import uk.ac.ebi.velvet.processor.StatisticsProcessor;

/**
 * @author mhaimel
 *
 */
public class VelvetExec {
	private static final int DEFAULT_CATEGORY = 1;
	private Log log = LogFactory.getLog(this.getClass());
	public static final String DEFAULT_NAME = "test";
	private static final String DEFAULT_EXEC_G = "velvetg";
	private static final String DEFAULT_EXEC_H = "velveth";

	private final File baseDir;
	private final Integer kmer;
	private String name;
	private String execH;
	private String execG;
	
	private SortedMap<Integer,Set<FileInfo>> cat2Files = new TreeMap<Integer,Set<FileInfo>>();
	private Map<Integer, Integer> cat2InslenMap = new HashMap<Integer, Integer>();
	private Map<Integer, Integer> cat2SdMap = new HashMap<Integer, Integer>();
	private Double kmerCov = null;
	private Double kmerCutoff = null;
	private int minPairCount = -1;
	private final AtomicBoolean force = new AtomicBoolean(false);
	
	public VelvetExec(File baseDir, Integer kmer) {
		this(baseDir,kmer,DEFAULT_NAME, DEFAULT_EXEC_H, DEFAULT_EXEC_G);
	}

	public VelvetExec(File baseDir, Integer kmer, String name, String execH, String execG) {
		this.baseDir = baseDir;
		this.kmer = kmer;
		this.name = name;
		this.execH = execH;
		this.execG = execG;
	}
	
	public void addFile(Integer cat, FileInfo info){
		if(isReadFile(info)){
			if(info.getFile().length() > 0){
				_addReadFile(cat, info);
			} else {
				getLog().warn("Ignore empty file " + info);
			}
		}
	}

	private boolean isReadFile(FileInfo info) {
		ReadType rType = info.getReadType();
		return rType.equals(ReadType.Short)
			|| rType.equals(ReadType.ShortPaired)
			|| rType.equals(ReadType.Long)
			|| rType.equals(ReadType.LongPaired);
	}

	private void _addReadFile(Integer cat, FileInfo info) {
		Set<FileInfo> set = getCat2Files().get(cat);
		if(null == set){
			set = new HashSet<FileInfo>();
			getCat2Files().put(cat, set);
		}
		set.add(info);
	}
	
	public void addInsertLength(Integer cat, Integer insLen){
		getCat2InslenMap().put(cat, insLen);
	}
	
	/**
	 * Runs {@link #runVelvetH()}, than {@link #runVelvetG()},
	 * process the intermediate statistics using {@link StatisticsProcessor}
	 * and reruns the {@link #runVelvetG(StatisticsInfo)} using
	 * the intermediate statistics information
	 * @return {@link StatisticsInfo} final statistics information
	 */
	public StatisticsInfo runVelvetAll(){
		StatisticsInfo stats = null;
		boolean isOk = runVelvetH();
		if(isOk){
			boolean rerun = runVelvetG(false);
			File statsFile = getStatsFile();
//			File contigsFile = getContigsFile();
			if(rerun && statsFile.exists()){
				
				StatisticsProcessor sProc = new StatisticsProcessor(statsFile);
				StatisticsInfo statsA = sProc.process();
				stats=statsA;
				runVelvetG(true);
				boolean rollback = false;
				if(statsFile.exists()){
					StatisticsInfo statsB = sProc.process();
					getLog().debug("Stats after Paire-end run " + getName() + ": " + stats);
					if(null == statsA.getN50()){
						// do nothing
					} else if(null == statsB.getN50()){
						rollback = true;
					} else if(statsB.getN50() < statsA.getN50()){
						rollback = true;
					}
				} else {
					rollback = true;
				}
				if(rollback){
					getLog().debug("Rollback to unpaired state: " + getName() + "!");
					runVelvetG(false);
				}
			}
		}
		return stats;
	}
	

	public StatisticsInfo rerunVelvet() {
		File statsFile = getStatsFile();
		if(!getGraph2File().exists()){
			boolean rerun = runVelvetG(false);
			if(!rerun){
				return null;
			}
			if(!statsFile.exists()){
				return null;
			}
		}
		runVelvetG(true);
		if(!statsFile.exists()){
			return null;
		}
		StatisticsProcessor sProc = new StatisticsProcessor(statsFile);
		StatisticsInfo info = sProc.process();
		if(info.getN50() == null){
			return null;
		}
		return info;
	}
	
	private boolean hasReadFiles() {
		Map<Integer, Set<FileInfo>> readFiles = getReadFiles();
		for(Entry<Integer, Set<FileInfo>> entry : readFiles.entrySet()){
			if(!entry.getValue().isEmpty()){
				return true;
			}
		}
		return false;
	}

	public boolean runVelvetG(boolean options) {
		List<String> args = new ArrayList<String>();
		if(options){
			if(minPairCount > 0){
				args.add("-min_pair_count");
				args.add(Integer.valueOf(minPairCount).toString());
			}
			args.add("-cov_cutoff");
			if(null != getExpectedKmerCutoff()){
				args.add(getExpectedKmerCutoff().toString());
			} else {
				args.add("auto");				
			}
			args.add("-exp_cov");
			if(null != getExpectedKmerCoverage()){
				args.add(getExpectedKmerCoverage().toString());
			} else {
				args.add("auto");
			}
			Integer offset = _getCategoryOffset();
			for(Entry<Integer, Integer> entry : getCat2InslenMap().entrySet()){
				// add category except default category
				Integer cat = entry.getKey() + offset;
				String insLenId = "-ins_length"+translateCategory(cat);
				args.add(insLenId);
				args.add(entry.getValue().toString());
			}			
			for(Entry<Integer, Integer> entry : getCat2SdMap().entrySet()){
				Integer cat = entry.getKey() + offset;
				// add category except default category
				String insSd = "-ins_length"+translateCategory(cat)+"_sd";
				args.add(insSd);
				args.add(entry.getValue().toString());
			}
		}
		return runVelvetG(args.toArray(new String[0]));
	}

	private File getNameDir(){
		return new File(this.baseDir,this.name);
	}
	
	private String getName() {
		return name;
	}

	private Log getLog() {
		return log;
	}

	public File getContigsFile(){
		return new File(getResultDir(),"contigs.fa");
	}

	public File getStatsFile() {
		return new File(getResultDir(),"stats.txt");
	}
	
	public File getGraph2File() {
		return new File(getResultDir(),"Graph2");
	}

	public boolean runVelvetG(String ... strArr) {		
		List<String> cmds = new ArrayList<String>();
		cmds.add(getExecG());
		cmds.add(getName());
		cmds.addAll(Arrays.asList(strArr));
		String[] cmdArr = cmds.toArray(new String[0]);
		ExecReport report = exec(cmdArr);
		boolean rerun = true;
		if(report.hasError()){
			rerun = false;
			boolean valid = _exceptableError(report);
			checkError(report, !valid);
		}
		return rerun;
	}

	public void setForce(boolean force){
		this.force.set(force);
	}

	@SuppressWarnings("unchecked")
	private boolean _exceptableError(ExecReport report) {
		if(this.force.get()){
			return true; // ignore any error
		}
		boolean valid = false;
		List<String> lines = Collections.emptyList();
		try {
			lines = FileUtils.readLines(report.err);
		} catch (IOException e) {
			throw new CurtainUncheckedException("Problem reading file " + report.err,e);
		}
		// has error output
		if(!lines.isEmpty()){
			String lastLine = lines.get(lines.size()-1);
			if(lastLine.contains("PreGraph file incomplete: No such file or directory")){
				getLog().debug("Velvetg Exceptable error found: >PreGraph file incomplete< for: " + report.getMsg());
				valid = true;
			} else if(lastLine.contains("velvetg: Graph file incomplete: Success")){
				getLog().debug("Velvetg Exceptable error found: >Graph file incomplete< for: " + report.getMsg());
				valid = true;
			}
		}
		return valid;
	}

	private String getExecG() {
		return execG;
	}

	public boolean runVelvetH() {
		if(!hasReadFiles()){
			getLog().warn("No Read files for " + getName() + " in " + this.getBaseDir());
			return false;
		}
		List<String> cmdList = new ArrayList<String>();
		cmdList.add(getExecH());
		cmdList.add(getName());
		cmdList.add(getKmer().toString());
		cmdList.addAll(getFileOptions());
		ExecReport report = exec(cmdList.toArray(new String[0]));
		checkError(report,true);
		return true;
	}

	private ExecReport exec(String[] cmdArr) {
		File logF = getNextVelvetLogFile();
		File errF = getNextVelvetErrFile();
		return exec(getBaseDir(),cmdArr,logF, errF);
	}

	private Collection<? extends String> getFileOptions() {
		Integer offSet = _getCategoryOffset();
		// Short , short_paired, long, long_paired
		List<String> args = new ArrayList<String>();
		ReadType[] paramOrder = ReadType.getVelvetParameterOrder();
		for(ReadType rt : paramOrder){
			for(Entry<Integer, Set<FileInfo>> entry : getReadFiles().entrySet()){
				for(FileInfo fi : entry.getValue()){
					if(fi.getFile().exists()){
						if(rt.equals(fi.getReadType())){
							String fTrans = translate(fi.getFileType());
							int category = entry.getKey() + offSet;
							String rTrans = translate(fi.getReadType(), category);
							String path = fi.getFile().getAbsolutePath();
							
							args.add(fTrans);
							args.add(rTrans);
							args.add(path);
						}
					}
				}
			}
		}
		return args;
	}
	
	private Integer _getCategoryOffset(){
		int minPaired = 0;
		myPoint: for(Entry<Integer, Set<FileInfo>> entry : getReadFiles().entrySet()){
			for(FileInfo fi : entry.getValue()){
				if(fi.getReadType().isPaired()){
					minPaired = entry.getKey();
					break myPoint;
				}
			}
		}
		return DEFAULT_CATEGORY - minPaired;
	}
	
	private String translate(FileType fileType) {
		if(FileType.fasta.equals(fileType)){
			return "-fasta";
		} else if(FileType.fastq.equals(fileType)){
			return "-fastq";
		} else {
			throw new CurtainUncheckedException("File type not supported! " + fileType); 
		}
	}
	
	private String translate(ReadType readType, Integer cat) {
		if(ReadType.ShortPaired.equals(readType)){
			return "-shortPaired"+translateCategory(cat);
		}
		if(ReadType.Short.equals(readType)){
			return "-short"+translateCategory(cat);
		}
		if(ReadType.Long.equals(readType)){
			return "-long";
		} 
		throw new CurtainUncheckedException("Read type not yet supported! " + readType); 
	}
	
	private String translateCategory(Integer cat) {
		if(cat > DEFAULT_CATEGORY){
			return cat.toString();
		} else {
			return StringUtils.EMPTY;
		}
	}
	
	public File getResultDir() {
		return new File(getBaseDir(),getName());
	}
	
	private File getBaseDir() {
		return baseDir;
	}
	
	public void addInsertLengthSd(Integer cat, Integer insSd) {
		cat2SdMap.put(cat, insSd);
	}
	
	private File getNextVelvetLogFile(){
		return getNextFile(getResultDir(),"velvet_", ".log");
	}
	private File getNextVelvetErrFile(){
		return getNextFile(getResultDir(),"velvet_", ".err");
	}
	
	protected File getNextFile(File dir, String prefix,String postfix){
		if(!dir.exists()){
			if(!dir.mkdirs()){
				throw new VelvetUncheckedException(
						"Not able to create File " + dir);
			}
		}
		int i = 0;
		File f = new File(dir,prefix+i+postfix);
		while(f.exists()){
			++i;
			f = new File(dir,prefix+i+postfix);
		}
		return f;
	}

	public Map<Integer, Set<FileInfo>> getReadFiles() {
		return getCat2Files();
	}

	protected String getExecH() {
		return execH;
	}
	
	protected Integer getKmer() {
		return kmer;
	}

	private Map<Integer, Integer> getCat2InslenMap() {
		return cat2InslenMap;
	}
	
	private Map<Integer, Integer> getCat2SdMap() {
		return cat2SdMap;
	}
	
	private Map<Integer, Set<FileInfo>> getCat2Files() {
		return cat2Files;
	}
	
	private ExecReport exec(File base, String[] cmdArr, File logFile, File errFile) {
		checkDiskDelay();
		ExecReport report = new ExecReport(logFile,errFile,-1);
		report.cmdArr = cmdArr;
		OutputStream logO = null;
		OutputStream errO = null;
		getLog().debug("Execute in " + base + "; cmd:[" + StringUtils.join(cmdArr, "; ")+"]");
		int exitCode = -1; 
		try {
			try{
				logO = new FileIO(logFile).getBufferedOutputStream();
				errO = new FileIO(errFile).getBufferedOutputStream();
				exitCode = buildRuntime(baseDir).execArr(
						logO,
						errO,
						cmdArr
						);
			}finally{
				IOUtils.closeQuietly(logO);
				IOUtils.closeQuietly(errO);
			}
			lastExec.set(System.currentTimeMillis());
			report.exit = exitCode;
			String msg = report.getMsg();
			getLog().debug("Finished " + msg);
		} catch (IOException e) {
			throw new VelvetUncheckedException("Problems executing command: " + StringUtils.join(cmdArr, "; "),e);
		}
		return report;
	}

	private final AtomicLong lastExec = new AtomicLong(0);
	private final AtomicLong diskDelay = new AtomicLong(Long.getLong("curtain.diskdelay", 50));
	private void checkDiskDelay() {
		long curr = System.currentTimeMillis();
		long diff = curr - lastExec.get();
		long delay = diskDelay.get();
		if(delay > diff){
			if(diff <= 0){
				diff = delay;
			} else {
				diff = delay -diff;
			}
			try {
				Thread.currentThread().sleep(diff);
			} catch (InterruptedException e) {
				Thread.currentThread().interrupt();
			}
		}
	}

	private void checkError(ExecReport report,boolean failOnError){
		String msg = report.getMsg();
		if(report.hasExitCodeError()){
			String log = "Problems with error code: "+ msg;
			if(failOnError){
				throw new VelvetUncheckedException(log);
			} else {
				getLog().warn(log);
			}
		}
		if(report.hasErrorLog()){
			String log = "Problems with Error output: "+ msg;
			if(failOnError){
				throw new VelvetUncheckedException(log);
			} else {
				getLog().warn(log);
			}		
		}
	}

	protected RuntimeExec buildRuntime(File baseDir) {
		return new RuntimeExec(baseDir,null);
	}

	private static class ExecReport{

		private File log;
		private File err;
		private int exit;	
		private String[] cmdArr;
		
		public ExecReport(File log, File err, int exit) {
			super();
			this.log = log;
			this.err = err;
			this.exit = exit;
		}
		
		private boolean hasExitCodeError(){
			return exit != 0;
		}
		private boolean hasErrorLog(){
			return err.length() > 0;
		}
		
		private boolean hasError(){
			if(hasExitCodeError()){
				return true;
			}
			if(hasErrorLog()){
				return true;
			}
			return false;
		}
		
		public String getMsg() {
			return "Exit code " + exit + " with cmd: "
			+ StringUtils.join(cmdArr, ',')
			+"! output:>"
			+log.getAbsolutePath()
			+"< err:>"
			+err.getAbsolutePath()
			+"<";
		}
	}

	public void setExpectedKmerCoverage(double kmerCov) {
		this.kmerCov  = kmerCov;
	}
	
	public Double getExpectedKmerCoverage() {
		return kmerCov;
	}

	public void setExpectedKmerCutoff(double kmerCutoff) {
		this.kmerCutoff  = kmerCutoff;
	}
	
	public Double getExpectedKmerCutoff() {
		return kmerCutoff;
	}

	public void setMinPairCount(int velvetMinPairCount) {
		this.minPairCount  = velvetMinPairCount;
	}

	
}
